"""
蚁群算法 非线性函数求极值
"""
import os

os.chdir(os.path.dirname(__file__))

import random
from copy import deepcopy
import math
import matplotlib.pyplot as plt #导入所需要的库

class ACO:
    def __init__(self, func, n = 300, var = 1, iter = 50, lb = None, ub = None):
        """ 默认寻找最大值，以及对应自变量取值
        :param func: 所要求解的目标函数
        :param n: 蚁群蚂蚁的个数
        :param var: 函数中自变量的个数，即：几元函数
        :param iter: 迭代次数
        :param lb: list(double)，每一种自变量的下界，顺序应和自变量一一对应
        :param ub: list(double)，每一种自变量的上界，顺序应和自变量一一对应
        """
        if lb is None:
            lb = []
        if ub is None:
            ub = []
        self.func = func  # 目标函数
        self.n = n        # 初始化蚂蚁数量
        self.var = var    # 初始化自变量个数
        self.iter = iter  # 初始化迭代次数
        self.lb = lb      # 初始化自变量范围 len(lb) = len(ub) = var
        self.ub = ub
        self.Rou = 0.9    # 初始化信息素挥发系数
        self.P0 = 0.8     # 初始化转移概率常数

        # 为每只蚂蚁的适应度申请空间
        ## 每只蚂蚁的信息素浓度。 取值与适应度（一般为函数值）有关
        self.phero = [0] * n
        # 为每只蚂蚁的坐标申请空间
        self.ant = [[0] * var for _ in range(n)]
        # 初始化每只蚂蚁的坐标
        for i in range(n):
            for j in range(var):
                self.ant[i][j] = lb[j] + random.random() * (ub[j] - lb[j]) # 使个自变量随机分布于可行域中
            self.phero[i] = self.func(self.ant[i])                         # 计算一组初始解
        self.max_x = deepcopy(self.ant[0])                    # 最优解位置
        # 初始化最优解位置
        for i in range(1, n):
            # 基于最大值选择最优值
            if self.phero[i] > self.func(self.max_x):
                self.max_x = deepcopy(self.ant[i])

    def solve(self):
        curr_max_x = []
        max_x = []
        for k in range(1, self.iter + 1):
            # 局部搜索步长
            ## 随迭代次数增加步长逐渐减小
            step = 1 / k

            # 每只蚂蚁的转移概率
            p = [0] * self.n
            ## 信息素越多， 转移概率越大
            for i in range(self.n):
                p[i] = self.phero[i] / max(self.phero)

            new_phero= [0] * self.n
            
            gen_max_x = self.ant[0]
            for i in range(self.n):
                for j in range(self.var):
                    if p[i] >= self.P0: # 局部搜索
                        # rand - 0.5 为了均匀产生正负位移
                        self.ant[i][j] += (random.random() - 0.5) * step
                    else: # 全局搜索
                        # rand - 0.5 为了均匀产生正负位移
                        self.ant[i][j] += (random.random() - 0.5) * (self.ub[j] - self.lb[j])
                    # 调整自变量值避免越界
                    if self.ant[i][j] < self.lb[j]:
                        self.ant[i][j] = self.lb[j]
                    if self.ant[i][j] > self.ub[j]:
                        self.ant[i][j] = self.ub[j]
                # 更新最优值位置
                if self.func(self.ant[i]) > self.func(gen_max_x):
                    gen_max_x = deepcopy(self.ant[i])
                new_phero[i] = self.func(self.ant[i])
            curr_max_x.append(gen_max_x)
            if  self.func(gen_max_x) > self.func(self.max_x):
                self.max_x = deepcopy(gen_max_x)
            max_x.append(self.max_x)
            for i in range(self.n):
                self.phero[i] = (1 - self.Rou) * self.phero[i] + new_phero[i]

        print(f'最优解为：{self.func(self.max_x)}')
        print(f'最优解对应自变量取值为：{self.max_x}')

        return self.max_x, max_x, curr_max_x



def func(x):
    # 目标函数
    return abs(x[0] * math.sin(x[0]) * math.cos(2 * x[0]) - 2 * x[0] * math.sin(3 * x[0]) + 3 * x[0] * math.sin(4 * x[0]))


aco = ACO(func, 50, 1, 50, [0], [50]) # x:0~50
max_x, max_x_log, curr_max_x_log = aco.solve()

n = len(max_x_log)
gen = []
max_y = []
curr_max_y = []
for i in range(n):
    gen.append(str(i+1))
    max_y.append(func(max_x_log[i]))
    curr_max_y.append(func(curr_max_x_log[i]))


# 适应值函数曲线
plt.figure(1)
plt.subplot(1,2,1)
plt.plot(gen, max_y, c='red', label="global best cost")
plt.plot(gen, curr_max_y, c='green', linestyle='--', label="generation best cost")
plt.legend(loc='best')

plt.figure(1)
plt.subplot(1,2,2)
plt.plot(gen, [i[0] for i in max_x_log], c='red', label="global best cost")
plt.plot(gen, [i[0] for i in curr_max_x_log], c='green', linestyle='--', label="generation best cost")
plt.legend(loc='best')

plt.show()